import cv2
import numpy as np
import matplotlib.pyplot as plt
import os

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 获取当前脚本所在目录
# script_dir = os.path.dirname(os.path.abspath(__file__))
# 图像路径 - 使用绝对路径
image_path = 'cat.jpg'

# 检查文件是否存在
print(f"检查图像文件: {image_path}")
if not os.path.exists(image_path):
    print(f"错误: 图像文件不存在: {image_path}")
    print("请确保在当前目录下放置一张名为'cat.jpg'的图像文件")
    # 创建一个示例图像用于演示
    print("创建示例图像用于演示...")
    # 创建一个彩色示例图像
    sample_img = np.random.randint(0, 255, (300, 300, 3), dtype=np.uint8)
    # 添加一些图形元素
    cv2.circle(sample_img, (150, 150), 50, (255, 0, 0), -1)  # 蓝色圆
    cv2.rectangle(sample_img, (50, 50), (250, 250), (0, 255, 0), 3)  # 绿色矩形
    cv2.putText(sample_img, 'Sample Image', (70, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)  # 红色文本
    cv2.imwrite(image_path, sample_img)
    print(f"示例图像已创建: {image_path}")

# 加载图像 (OpenCV默认加载为BGR格式，需要转换为RGB)
print(f"尝试加载图像: {image_path}")
img = cv2.imread(image_path)

# 检查图像是否成功加载
if img is None:
    raise RuntimeError(f"无法加载图像文件: {image_path}\n"
                      f"请检查文件路径是否正确，文件是否损坏，以及是否为支持的图像格式")

print(f"图像加载成功，形状: {img.shape}")

# 转换为RGB格式
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换为RGB格式
img_array = np.array(img)

# 定义图像增强函数

def random_flip(image):
    """随机水平或垂直翻转"""
    # 50%的概率水平翻转
    if np.random.random() > 0.5:
        image = cv2.flip(image, 1)
    # 50%的概率垂直翻转
    if np.random.random() > 0.5:
        image = cv2.flip(image, 0)
    return image

def random_rotation(image, max_angle=30):
    """随机旋转图像"""
    angle = np.random.uniform(-max_angle, max_angle)
    height, width = image.shape[:2]
    center = (width // 2, height // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_image = cv2.warpAffine(image, rotation_matrix, (width, height))
    return rotated_image

def random_zoom(image, min_zoom=0.8, max_zoom=1.2):
    """随机缩放图像"""
    zoom_factor = np.random.uniform(min_zoom, max_zoom)
    height, width = image.shape[:2]
    new_height = int(height * zoom_factor)
    new_width = int(width * zoom_factor)
    # 缩放图像
    image = cv2.resize(image, (new_width, new_height))
    # 裁剪或填充回原始大小
    if zoom_factor > 1:
        # 裁剪
        start_x = (new_width - width) // 2
        start_y = (new_height - height) // 2
        image = image[start_y:start_y+height, start_x:start_x+width]
    else:
        # 填充
        pad_top = (height - new_height) // 2
        pad_bottom = height - new_height - pad_top
        pad_left = (width - new_width) // 2
        pad_right = width - new_width - pad_left
        image = cv2.copyMakeBorder(image, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
    return image

def random_brightness(image, max_delta=50):
    """随机调整亮度"""
    delta = np.random.uniform(-max_delta, max_delta)
    image = np.clip(image + delta, 0, 255).astype(np.uint8)
    return image

def random_contrast(image, alpha_range=(0.8, 1.2)):
    """随机调整对比度"""
    alpha = np.random.uniform(alpha_range[0], alpha_range[1])
    image = np.clip(alpha * image, 0, 255).astype(np.uint8)
    return image

def random_saturation(image, saturation_range=(0.5, 1.5)):
    """随机调整饱和度"""
    # 转换到HSV色彩空间
    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    # 调整饱和度
    saturation_factor = np.random.uniform(saturation_range[0], saturation_range[1])
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * saturation_factor, 0, 255)
    # 转换回RGB
    image = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
    return image

# 应用增强并显示结果

# 创建一个图像增强管道
def augment_image(image):
    image = random_flip(image)
    image = random_rotation(image)
    image = random_zoom(image)
    image = random_brightness(image)
    image = random_contrast(image)
    image = random_saturation(image)
    return image

# 生成多个增强样本
num_samples = 6
augmented_images = [augment_image(img_array.copy()) for _ in range(num_samples)]

# 显示原始图像和增强后的图像
plt.figure(figsize=(15, 10))

# 显示原始图像
plt.subplot(2, 4, 1)
plt.imshow(img_array.astype(np.uint8))
plt.title("原始图像")
plt.axis('off')

# 显示增强后的图像
for i, aug_img in enumerate(augmented_images):
    plt.subplot(2, 4, i+2)
    plt.imshow(aug_img.astype(np.uint8))
    plt.title(f"增强图像 {i+1}")
    plt.axis('off')

plt.tight_layout()
plt.show()

print("图像增强演示完成！")